import numpy as np
import torch
from collections import defaultdict
from gensim.models import Word2Vec

def load_graph_data(graph_name):
    """加载PyTorch格式的图数据"""
    data = torch.load(f"../data/{graph_name}_data.pt")[0]  # 假设数据存储在tuple中
    edge_index = data['edge_index'].numpy().T  # 转换为Nx2的边数组
    num_nodes = data['x'].shape[0]  # 通过特征矩阵维度获取节点总数
    return edge_index, num_nodes

def build_adjacency(edge_index, num_nodes):
    """构建邻接表并确保所有节点存在"""
    adj_list = defaultdict(list)
    nodes = set()
    
    # 添加双向边（假设是无向图）
    for u, v in edge_index:
        adj_list[str(u)].append(str(v))
        adj_list[str(v)].append(str(u))
        nodes.update([str(u), str(v)])
    
    # 补全孤立节点
    for node in range(num_nodes):
        str_node = str(node)
        if str_node not in adj_list:
            adj_list[str_node] = []
    
    return adj_list

def generate_sentences(adj_list, num_nodes):
    """生成训练序列并确保节点顺序"""
    sentences = []
    for node_id in range(num_nodes):
        node_str = str(node_id)
        neighbors = adj_list[node_str]
        if len(neighbors) == 0:
            sentences.append([node_str])
        else:
            sentences.append([node_str] + neighbors)
    return sentences

def train_line(sentences, emb_dim=128):
    """训练LINE模型（二阶相似性）"""
    model = Word2Vec(
        sentences,
        vector_size=emb_dim,
        window=5,
        sg=1,
        hs=1,
        min_count=0,
        workers=4,
        epochs=10
    )
    return model

def save_embeddings(model, num_nodes, emb_dim, output_file):
    """保存符合格式要求的emb文件"""
    with open(output_file, 'w') as f:
        f.write(f"{num_nodes}\t{emb_dim}\n")
        for node_id in range(num_nodes):
            node_str = str(node_id)
            vector = model.wv[node_str] if node_str in model.wv else np.zeros(emb_dim)
            line = [node_str] + [f"{x:.6f}" for x in vector]
            f.write("\t".join(line) + "\n")

if __name__ == "__main__":
    # 用户输入图名称（如Cora/Citeseer）
    graph_name = input("请输入图名称（例如Cora/Citeseer）: ").strip()
    
    # 参数设置
    emb_dim = 128
    output_file = f"../data/{graph_name}.emb"
    
    # 处理流程
    edge_index, num_nodes = load_graph_data(graph_name)
    adj_list = build_adjacency(edge_index, num_nodes)
    sentences = generate_sentences(adj_list, num_nodes)
    model = train_line(sentences, emb_dim)
    save_embeddings(model, num_nodes, emb_dim, output_file)
    print(f"嵌入文件已生成：{output_file}")